import numpy as np
from typing import List, Tuple

class RobotArm:
    def __init__(self):
        # DH参数初始化，根据表1
        self.dh_params = {
            'a': [0, 500, 0, 0, 0, 0],  # 连杆长度(mm)
            'alpha': [np.pi/2, 0, np.pi/2, -np.pi/2, np.pi/2, 0],  # 扭角
            'd': [400, 0, 0, 500, 0, 100],  # 偏距
            'theta': [0, -np.pi/2, 0, np.pi, -np.pi/2, 0]  # 初始关节角度
        }
        
        # 关节限制（更新为更合理的范围）
        self.joint_limits = {
            'lower': [-170, -170, -170, -170, -170, -170],  # 度
            'upper': [170, 170, 170, 170, 170, 170]  # 度
        }
        
    def transform_matrix(self, theta: float, d: float, a: float, alpha: float) -> np.ndarray:
        """计算单个关节的变换矩阵"""
        ct = np.cos(theta)
        st = np.sin(theta)
        ca = np.cos(alpha)
        sa = np.sin(alpha)
        
        return np.array([
            [ct, -st*ca, st*sa, a*ct],
            [st, ct*ca, -ct*sa, a*st],
            [0, sa, ca, d],
            [0, 0, 0, 1]
        ])
    
    def forward_kinematics(self, joint_angles: List[float]) -> np.ndarray:
        """前向运动学计算"""
        T = np.eye(4)
        for i in range(6):
            theta = joint_angles[i] + self.dh_params['theta'][i]
            d = self.dh_params['d'][i]
            a = self.dh_params['a'][i]
            alpha = self.dh_params['alpha'][i]
            
            Ti = self.transform_matrix(theta, d, a, alpha)
            T = T @ Ti
            
        return T
    
    def calculate_end_effector_position(self, joint_angles: List[float]) -> np.ndarray:
        """计算末端执行器位置"""
        T = self.forward_kinematics(joint_angles)
        return T[:3, 3]
    
    def calculate_end_effector_error(self, joint_angles: List[float], target_position: np.ndarray) -> float:
        """计算末端误差"""
        current_position = self.calculate_end_effector_position(joint_angles)
        return np.linalg.norm(current_position - target_position)
    
    def get_random_angles(self) -> List[float]:
        """生成随机的关节角度（用于优化初始值）"""
        random_angles = []
        for i in range(6):
            lower = np.deg2rad(self.joint_limits['lower'][i])
            upper = np.deg2rad(self.joint_limits['upper'][i])
            random_angles.append(np.random.uniform(lower, upper))
        return random_angles
        
    def calculate_joint_positions(self, joint_angles: List[float]) -> List[np.ndarray]:
        """
        计算所有关节的位置
        
        参数:
            joint_angles: 关节角度列表
            
        返回:
            所有关节位置的列表，包括基座(关节0)和末端执行器
        """
        positions = []
        T = np.eye(4)  # 初始变换矩阵（基座位置）
        
        # 添加基座位置
        positions.append(T[:3, 3].copy())
        
        # 逐个计算每个关节的位置
        for i in range(6):
            theta = joint_angles[i] + self.dh_params['theta'][i]
            d = self.dh_params['d'][i]
            a = self.dh_params['a'][i]
            alpha = self.dh_params['alpha'][i]
            
            Ti = self.transform_matrix(theta, d, a, alpha)
            T = T @ Ti
            
            # 储存当前关节位置
            positions.append(T[:3, 3].copy())
        
        return positions
    
    def check_collision(self, joint_positions: List[np.ndarray], obstacles: List[np.ndarray], safety_margin: float = 100) -> bool:
        """
        检查机器臂是否与障碍物碰撞
        
        参数:
            joint_positions: 关节位置列表
            obstacles: 障碍物位置列表，每个元素为[x, y, z]坐标
            safety_margin: 安全边际(mm)
            
        返回:
            如果有碰撞则返回True，否则返回False
        """
        # 检查每个关节与每个障碍物之间的距离
        for joint_pos in joint_positions:
            for obstacle in obstacles:
                # 将障碍物坐标转换为3D，假设障碍物高度为0-1000mm
                obs_3d = np.array([obstacle[0], obstacle[1], 0])
                
                # 计算关节到障碍物的水平距离（忽略z轴）
                horizontal_distance = np.linalg.norm(joint_pos[:2] - obs_3d[:2])
                
                # 如果水平距离小于安全边际，并且关节z坐标小于障碍物高度上限，则视为碰撞
                if horizontal_distance < safety_margin and joint_pos[2] < 1000:
                    return True
        
        # 检查连杆与障碍物之间的碰撞
        for i in range(len(joint_positions) - 1):
            joint1 = joint_positions[i]
            joint2 = joint_positions[i + 1]
            
            # 连杆方向向量
            link_vector = joint2 - joint1
            link_length = np.linalg.norm(link_vector)
            
            if link_length > 0:
                # 检查连杆上的多个点
                num_checks = 10  # 检查连杆上的点数
                for t in np.linspace(0, 1, num_checks):
                    point = joint1 + t * link_vector
                    
                    for obstacle in obstacles:
                        obs_3d = np.array([obstacle[0], obstacle[1], 0])
                        horizontal_distance = np.linalg.norm(point[:2] - obs_3d[:2])
                        
                        if horizontal_distance < safety_margin and point[2] < 1000:
                            return True
        
        return False